#!/usr/bin/env python3

import numpy as np
import cv2
import depthai as dai
import time
from datetime import timedelta
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from builtin_interfaces.msg import Time
from cv_bridge import CvBridge

# 配置参数
FPS = 30.0
RGB_SOCKET = dai.CameraBoardSocket.CAM_A
LEFT_SOCKET = dai.CameraBoardSocket.CAM_B
RIGHT_SOCKET = dai.CameraBoardSocket.CAM_C
ALIGN_SOCKET = LEFT_SOCKET
COLOR_RESOLUTION = dai.ColorCameraProperties.SensorResolution.THE_1080_P
LEFT_RIGHT_RESOLUTION = dai.MonoCameraProperties.SensorResolution.THE_400_P
ISP_SCALE = 3

class DepthAIRosPublisher(Node):
    def __init__(self):
        super().__init__('depthai_ros_publisher')
        self.bridge = CvBridge()
        
        # 创建发布者
        self.rgb_pub = self.create_publisher(
            Image, 
            '/camera/aligned_rgb', 
            10
        )
        self.depth_pub = self.create_publisher(
            Image, 
            '/camera/aligned_depth', 
            10
        )
        
        # 初始化设备
        self.device = dai.Device()
        self.calibration = self.device.readCalibration()
        
        # 创建pipeline
        self.pipeline = self.create_pipeline()
        self.device.startPipeline(self.pipeline)
        self.queue = self.device.getOutputQueue("out", 8, False)
        
        # 帧ID和参数
        self.frame_id = "oak_pro_frame"
        self.rgb_distortion = self.calibration.getDistortionCoefficients(RGB_SOCKET)
        self.rgb_intrinsics = None  # 将在运行时动态获取

    def create_pipeline(self):
        pipeline = dai.Pipeline()

        # 定义节点
        cam_rgb = pipeline.create(dai.node.ColorCamera)
        left = pipeline.create(dai.node.MonoCamera)
        right = pipeline.create(dai.node.MonoCamera)
        stereo = pipeline.create(dai.node.StereoDepth)
        align = pipeline.create(dai.node.ImageAlign)
        sync = pipeline.create(dai.node.Sync)
        xout = pipeline.create(dai.node.XLinkOut)

        # 配置RGB相机
        cam_rgb.setBoardSocket(RGB_SOCKET)
        cam_rgb.setResolution(COLOR_RESOLUTION)
        cam_rgb.setFps(FPS)
        cam_rgb.setIspScale(1, ISP_SCALE)

        # 配置单目相机
        left.setResolution(LEFT_RIGHT_RESOLUTION)
        left.setBoardSocket(LEFT_SOCKET)
        left.setFps(FPS)
        right.setResolution(LEFT_RIGHT_RESOLUTION)
        right.setBoardSocket(RIGHT_SOCKET)
        right.setFps(FPS)

        # 配置立体深度
        stereo.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.HIGH_DENSITY)
        stereo.setDepthAlign(ALIGN_SOCKET)
        stereo.setOutputSize(*cam_rgb.getIspSize())

        # 配置同步节点
        sync.setSyncThreshold(timedelta(seconds=0.5 / FPS))
        xout.setStreamName("out")

        # 连接节点
        cam_rgb.isp.link(sync.inputs["rgb"])
        left.out.link(stereo.left)
        right.out.link(stereo.right)
        stereo.depth.link(align.input)
        align.outputAligned.link(sync.inputs["depth_aligned"])
        cam_rgb.isp.link(align.inputAlignTo)
        sync.out.link(xout.input)

        return pipeline

    def get_ros_time(self, dai_timestamp):
        """转换DepthAI时间戳到ROS时间格式"""
        return Time(
            sec=int(dai_timestamp.seconds),
            nanosec=int(dai_timestamp.microseconds * 1000)
        )

    def publish_frames(self):
        while rclpy.ok():
            message_group = self.queue.get()
            
            # 获取原始帧数据
            rgb_frame = message_group["rgb"]
            depth_frame = message_group["depth_aligned"]

            # 处理RGB帧
            cv_rgb = rgb_frame.getCvFrame()
            
            # 动态获取当前分辨率下的内参
            if self.rgb_intrinsics is None or \
               self.rgb_intrinsics[0] != cv_rgb.shape[1] or \
               self.rgb_intrinsics[1] != cv_rgb.shape[0]:
                self.rgb_intrinsics = self.calibration.getCameraIntrinsics(
                    RGB_SOCKET, 
                    cv_rgb.shape[1], 
                    cv_rgb.shape[0]
                )
            
            # 去畸变
            cv_rgb_undistorted = cv2.undistort(
                cv_rgb,
                np.array(self.rgb_intrinsics),
                np.array(self.rgb_distortion)
            )

            # 处理深度帧（保持原始数据）
            depth_data = depth_frame.getFrame()  # 原始16位数据（单位：毫米）

            # 创建ROS消息
            timestamp = self.get_ros_time(rgb_frame.getTimestamp())
            
            # 发布RGB
            rgb_msg = self.bridge.cv2_to_imgmsg(cv_rgb_undistorted, encoding="bgr8")
            rgb_msg.header.stamp = timestamp
            rgb_msg.header.frame_id = self.frame_id
            self.rgb_pub.publish(rgb_msg)

            # 发布深度（16UC1格式，单位毫米）
            depth_msg = self.bridge.cv2_to_imgmsg(depth_data, encoding="16UC1")
            depth_msg.header.stamp = timestamp
            depth_msg.header.frame_id = self.frame_id
            self.depth_pub.publish(depth_msg)

            # 可选：显示图像用于调试
            cv2.imshow("Aligned RGB", cv_rgb_undistorted)
            
            # 将深度数据归一化用于显示（不影响发布的原始数据）
            depth_vis = (depth_data * (255.0 / 5000.0)).astype(np.uint8)  # 假设5米为最大可视范围
            cv2.imshow("Aligned Depth", depth_vis)
            
            if cv2.waitKey(1) == ord('q'):
                break

def main(args=None):
    rclpy.init(args=args)
    
    try:
        publisher = DepthAIRosPublisher()
        publisher.publish_frames()
    except KeyboardInterrupt:
        pass
    finally:
        # 清理资源
        publisher.destroy_node()
        rclpy.shutdown()
        cv2.destroyAllWindows()

if __name__ == '__main__':
    main()
